import tensorflow as tf
import os
import matplotlib.pyplot as plt

tf.random.set_seed(777)

root_path = '../../../../large_data/DL2/_many_files/cifar2_fast/'
train_path = root_path + 'train/*/*.jpg'
test_path = root_path + 'test/*/*.jpg'


def load_img(path):
    label = tf.where(tf.strings.regex_full_match(path, r'.*[\\/]airplane[\\/][^\\/]+$'), 1, 0)

    py_path = bytes.decode(path.numpy())
    _, filename = os.path.split(py_path)

    img = tf.io.read_file(path)
    img = tf.image.decode_jpeg(img)
    img = tf.image.resize(img, (32, 32)) / 255.
    return (img, label, filename)


def get_ds(path, batch_size):
    ds = tf.data.Dataset.list_files(path)\
        .map(lambda x: tf.py_function(load_img, [x], [tf.float32, tf.int32, tf.string]), num_parallel_calls=tf.data.experimental.AUTOTUNE)\
        .shuffle(buffer_size=1000)\
        .batch(batch_size)\
        .prefetch(tf.data.experimental.AUTOTUNE)
    return ds


batch_size = 8
ds_train = get_ds(train_path, batch_size)
ds_test = get_ds(test_path, batch_size)

x = []
y = []
name = []
for bx, by, bname in ds_train:
    x.append(bx)
    y.append(by)
    name.append(bname)

x = tf.concat(x, axis=0)
y = tf.concat(y, axis=0)
name = tf.concat(name, axis=0)
print(tf.shape(x))
print(tf.shape(y))
print(tf.shape(name))

plt.figure(figsize=[12, 6])
spr = 4
spc = 8
spn = 0
for i, xi in enumerate(x):
    if i >= spr * spc:
        break
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.imshow(xi)
    plt.axis('off')
    plt.title(bytes.decode(name[i].numpy()) + ': ' + str(y[i].numpy()))
